import os
import uuid
import base64
import grpc
from concurrent import futures
import image_service_pb2
import image_service_pb2_grpc
import math
import cv2
import numpy as np
from scipy.ndimage import gaussian_filter

# 图像质量评估类，包含计算MSE、PSNR和SSIM的方法
class ImageQualityMetrics:
    @staticmethod
    def calculate_mse(img1, img2):
        """计算均方误差 (MSE)"""
        return np.mean((img1 - img2) ** 2)

    @staticmethod
    def calculate_psnr(img1, img2):
        """计算峰值信噪比 (PSNR)"""
        mse = ImageQualityMetrics.calculate_mse(img1, img2)
        if mse == 0:
            return float('inf')
        max_pixel = 255.0
        return 20 * math.log10(max_pixel / math.sqrt(mse))

    @staticmethod
    def calculate_ssim(img1, img2):
        """计算结构相似性指数 (SSIM) - 使用OpenCV实现"""
        if len(img1.shape) == 3:
            img1 = cv2.cvtColor(img1, cv2.COLOR_BGR2GRAY)
        if len(img2.shape) == 3:
            img2 = cv2.cvtColor(img2, cv2.COLOR_BGR2GRAY)

        C1 = (0.01 * 255) ** 2
        C2 = (0.03 * 255) ** 2

        img1 = img1.astype(np.float64)
        img2 = img2.astype(np.float64)
        kernel = cv2.getGaussianKernel(11, 1.5)
        window = np.outer(kernel, kernel.transpose())

        mu1 = cv2.filter2D(img1, -1, window)[5:-5, 5:-5]
        mu2 = cv2.filter2D(img2, -1, window)[5:-5, 5:-5]
        mu1_sq = mu1 ** 2
        mu2_sq = mu2 ** 2
        mu1_mu2 = mu1 * mu2
        sigma1_sq = cv2.filter2D(img1 ** 2, -1, window)[5:-5, 5:-5] - mu1_sq
        sigma2_sq = cv2.filter2D(img2 ** 2, -1, window)[5:-5, 5:-5] - mu2_sq
        sigma12 = cv2.filter2D(img1 * img2, -1, window)[5:-5, 5:-5] - mu1_mu2

        ssim_map = ((2 * mu1_mu2 + C1) * (2 * sigma12 + C2)) / ((mu1_sq + mu2_sq + C1) * (sigma1_sq + sigma2_sq + C2))
        return ssim_map.mean()

# 超分辨率处理系统类，包含预处理、超分辨率重建和后处理的方法
class SuperResolutionSystem:
    """超分辨率处理系统"""

    def __init__(self):
        self.scale_factor = 2
        self.denoise_level = 1.0
        self.sharpen_level = 1.0
        self.sharpen_kernel_size = 3

    def preprocess(self, img):
        """图像预处理"""
        img = cv2.GaussianBlur(img, (5, 5), self.denoise_level)

        if len(img.shape) == 2:
            img = cv2.equalizeHist(img)
        else:
            img_yuv = cv2.cvtColor(img, cv2.COLOR_BGR2YUV)
            clahe = cv2.createCLAHE(clipLimit=2.0, tileGridSize=(8, 8))
            img_yuv[:, :, 0] = clahe.apply(img_yuv[:, :, 0])
            img = cv2.cvtColor(img_yuv, cv2.COLOR_YUV2BGR)
        return img

    def super_resolution(self, img, method='bicubic'):
        """超分辨率重建"""
        h, w = img.shape[:2]
        new_h, new_w = h * self.scale_factor, w * self.scale_factor

        if method == 'nearest':
            return cv2.resize(img, (new_w, new_h), interpolation=cv2.INTER_NEAREST)
        elif method == 'bilinear':
            return cv2.resize(img, (new_w, new_h), interpolation=cv2.INTER_LINEAR)
        elif method == 'bicubic':
            return cv2.resize(img, (new_w, new_h), interpolation=cv2.INTER_CUBIC)
        elif method == 'lanczos':
            return cv2.resize(img, (new_w, new_h), interpolation=cv2.INTER_LANCZOS4)
        else:
            raise ValueError(f"Unknown super-resolution method: {method}")

    def postprocess(self, img):
        """后处理增强"""
        img = self.adaptive_sharpen(
            img,
            kernel_size=self.sharpen_kernel_size,
            strength=self.sharpen_level
        )

        img = cv2.fastNlMeansDenoisingColored(
            img,
            None,
            h=int(self.denoise_level * 10),
            templateWindowSize=7,
            searchWindowSize=21
        )
        return img

    def adaptive_sharpen(self, img, kernel_size=3, strength=1.0):
        """自适应锐化方法"""
        img_float = img.astype(np.float32)

        lab = cv2.cvtColor(img, cv2.COLOR_BGR2LAB)
        l_channel = lab[:, :, 0].astype(np.float32)
        blur = cv2.GaussianBlur(l_channel, (15, 15), 0)
        local_contrast = np.abs(l_channel - blur)
        contrast_mask = cv2.normalize(local_contrast, None, 0.0, 1.0, cv2.NORM_MINMAX, dtype=cv2.CV_32F)

        kernel = np.array([[-1, -1, -1],
                           [-1, 8 + strength, -1],
                           [-1, -1, -1]], dtype=np.float32)
        kernel /= kernel.sum()

        sharpened = cv2.filter2D(img_float, -1, kernel)

        result = np.zeros_like(img_float)
        for c in range(3):
            result[:, :, c] = (
                    img_float[:, :, c] * (1.0 - contrast_mask) +
                    sharpened[:, :, c] * contrast_mask
            )

        return np.clip(result, 0, 255).astype(np.uint8)

# 图像服务类，实现gRPC接口
class ImageService(image_service_pb2_grpc.ImageServiceServicer):
    def __init__(self):
        self.sr_system = SuperResolutionSystem()

    def UploadImage(self, request, context):
        try:
            # 解码文件内容
            file_data = base64.b64decode(request.file)

            # 获取文件扩展名
            original_filename = request.file.split(",")[0]+".jpg"
            _, ext = os.path.splitext(original_filename)
            ext = ext.lower()

            # 生成新的文件名
            new_name = str(uuid.uuid4()).replace("-", "") + ext+".jpg"

            # 原始图片路径
            original_path = os.path.join("original", new_name)  # 修改为保存到 original 文件夹
            os.makedirs(os.path.dirname(original_path), exist_ok=True)

            # 保存原始图片
            with open(original_path, "wb") as f:
                f.write(file_data)

            # 读取图片并进行处理
            img = cv2.imdecode(np.frombuffer(file_data, np.uint8), cv2.IMREAD_COLOR)
            if img is None:
                raise ValueError("无法解码图像")

            # 预处理
            preprocessed = self.sr_system.preprocess(img)

            # 超分辨率重建
            sr_result = self.sr_system.super_resolution(preprocessed, method='bicubic')

            # 后处理
            final_result = self.sr_system.postprocess(sr_result)

            # 处理后的图片路径
            display_path = os.path.join("display", new_name)  # 修改为保存到 display 文件夹
            os.makedirs(os.path.dirname(display_path), exist_ok=True)

            # 保存处理后的图片
            cv2.imwrite(display_path, final_result)

            # 构造返回信息
            response = image_service_pb2.UploadResponse(
                code=200,
                data=image_service_pb2.ImageInfo(
                    name=original_filename,
                    url=original_path.replace("\\", "/"),
                    newname=new_name,
                    newurl=display_path.replace("\\", "/")
                ),
                msg="上传和处理成功"
            )
            return response
        except Exception as e:
            return image_service_pb2.UploadResponse(code=500, msg=str(e))

# 启动gRPC服务器
def serve():
    server = grpc.server(futures.ThreadPoolExecutor(max_workers=10))
    image_service_pb2_grpc.add_ImageServiceServicer_to_server(ImageService(), server)
    server.add_insecure_port('[::]:50051')
    print("Python RPC 服务已启动在 50051 端口")

    server.start()
    server.wait_for_termination()

# 主入口
if __name__ == '__main__':
    serve()
